import { btPrint } from "hy-algokit";

enum Color {
  RED,
  BLACK,
}

class RedBlackNode<T> {
  value: T;
  color: Color;
  parent: RedBlackNode<T> | null;
  left: RedBlackNode<T> | null;
  right: RedBlackNode<T> | null;

  constructor(
    value: T,
    color: Color = Color.RED,
    parent: RedBlackNode<T> | null = null,
    left: RedBlackNode<T> | null = null,
    right: RedBlackNode<T> | null = null
  ) {
    this.value = value;
    this.color = color;
    this.parent = parent;
    this.left = left;
    this.right = right;
  }
}

class RedBlackTree<T> {
  root: RedBlackNode<T> | null = null;

  /**
   * 插入一个新节点
   * @param value 待插入节点的值
   */
  insert(value: T) {
    // 创建一个新节点
    let newNode = new RedBlackNode(value);

    // 如果红黑树为空，将该节点作为根节点
    if (!this.root) {
      this.root = newNode;
      // 根节点为黑色
      newNode.color = Color.BLACK;
      return;
    }

    // 初始化搜索变量current和parent
    let current: RedBlackNode<T> | null = this.root;
    let parent: RedBlackNode<T> | null = null;

    // 搜索合适的插入位置
    while (current) {
      parent = current;
      // 如果value小于当前节点，则继续往左子树搜索
      if (value < current.value) {
        current = current.left;
        // 否则继续往右子树搜索
      } else {
        current = current.right;
      }
    }

    // 将新节点的父节点设置为搜索到的父节点
    newNode.parent = parent;
    // 将新节点插入到合适的位置
    if (value < parent!.value) {
      parent!.left = newNode;
    } else {
      parent!.right = newNode;
    }

    // 修复插入导致的红黑树性质破坏
    this.fixInsertion(newNode);
  }

  private fixInsertion(node: RedBlackNode<T>) {
    // 当父节点存在且颜色为红时
    while (node.parent && node.parent.color === Color.RED) {
      // 获取祖父节点
      let grandParent = node.parent.parent!;

      // 父节点是祖父节点的左子节点
      if (node.parent === grandParent.left) {
        // 获取叔叔节点
        let uncle = grandParent.right;
        // 叔叔节点存在且颜色为红
        if (uncle && uncle.color === Color.RED) {
          // 将父节点颜色改为黑，叔叔节点颜色改为黑，祖父节点颜色改为红，node节点变为祖父节点，继续循环
          node.parent.color = Color.BLACK;
          uncle.color = Color.BLACK;
          grandParent.color = Color.RED;
          node = grandParent;
        } else {
          // 当前节点是父节点的右子节点
          if (node === node.parent.right) {
            // 将当前节点变为父节点，进行左旋操作
            node = node.parent;
            this.leftRotate(node);
          }
          // 将父节点颜色改为黑，祖父节点颜色改为红，进行右旋操作
          node.parent!.color = Color.BLACK;
          grandParent.color = Color.RED;
          this.rightRotate(grandParent);
        }
      } else {
        // 父节点是祖父节点的右子节点，与上面的同理
        let uncle = grandParent.left;
        // 如果叔叔节点是红色的
        if (uncle && uncle.color === Color.RED) {
          // 父节点设置为黑色
          node.parent.color = Color.BLACK;
          // 叔叔节点设置为黑色
          uncle.color = Color.BLACK;
          // 祖父节点设置为红色
          grandParent.color = Color.RED;
          // 将当前节点设置为祖父节点
          node = grandParent;
        } else {
          // 如果当前节点是父节点的左节点
          if (node === node.parent.left) {
            // 将当前节点设置为父节点
            node = node.parent;
            // 右旋父节点
            this.rightRotate(node);
          }
          // 父节点设置为黑色
          node.parent!.color = Color.BLACK;
          // 祖父节点设置为红色
          grandParent.color = Color.RED;
          // 左旋祖父节点
          this.leftRotate(grandParent);
        }
      }
    }
    // 根节点设置为黑色节点
    this.root!.color = Color.BLACK;
  }

  /**
   * 左旋操作
   *
   * @param node 要进行左旋的结点
   */
  private leftRotate(node: RedBlackNode<T>) {
    // 获取 node 的右子节点
    let rightChild = node.right!;
    // 将右子节点的左子节点赋值给 node 的右子节点
    node.right = rightChild.left;

    // 如果右子节点的左子节点不为空，则将右子节点的左子节点的父节点指向 node
    if (rightChild.left) {
      rightChild.left.parent = node;
    }

    // 将右子节点的父节点指向 node 的父节点
    rightChild.parent = node.parent;
    // 如果 node 的父节点为空，则将右子节点设为根结点
    if (!node.parent) {
      this.root = rightChild;
    }
    // 如果 node 是它父节点的左子节点，则将右子节点设为 node 父节点的左子节点
    else if (node === node.parent.left) {
      node.parent.left = rightChild;
    }
    // 否则，将右子节点设为 node 父节点的右子节点
    else {
      node.parent.right = rightChild;
    }

    // 将 node 的父节点指向 rightChild，并将 rightChild 的左子节点指向 node
    rightChild.left = node;
    node.parent = rightChild;
  }

  /**
   * 右旋转
   * @param node 旋转节点
   */
  private rightRotate(node: RedBlackNode<T>) {
    // 获取旋转节点的左子节点
    let leftChild = node.left!;
    // 将旋转节点的左子节点的右子节点，接到旋转节点的左边
    node.left = leftChild.right;

    // 如果左子节点的右子节点不为空，设置它的父节点为旋转节点
    if (leftChild.right) {
      leftChild.right.parent = node;
    }

    // 将左子节点的父节点设为旋转节点的父节点
    leftChild.parent = node.parent;
    // 如果旋转节点的父节点不存在，说明左子节点变成根节点
    if (!node.parent) {
      this.root = leftChild;
    } else if (node === node.parent.right) {
      // 如果旋转节点是它父节点的右子节点，将父节点的右子节点设为左子节点
      node.parent.right = leftChild;
    } else {
      // 如果旋转节点是它父节点的左子节点，将父节点的左子节点设为左子节点
      node.parent.left = leftChild;
    }

    // 将旋转节点设为左子节点的右子节点
    leftChild.right = node;
    // 将旋转节点的父节点设为左子节点
    node.parent = leftChild;
  }

  // 查找红黑树中的最小值
  minimum(node: RedBlackNode<T> | null = this.root): RedBlackNode<T> | null {
    let current = node;
    while (current && current.left) {
      current = current.left;
    }
    return current;
  }

  // 查找红黑树中的某个节点
  private search(value: T): RedBlackNode<T> | null {
    let node = this.root;
    let parent: RedBlackNode<T> | null = null;
    while (node) {
      if (node.value === value) {
        node.parent = parent;
        return node;
      }
      parent = node;
      if (value < node.value) {
        node = node.left;
      } else {
        node = node.right;
      }
    }
    return null;
  }

  /**
   * 删除红黑树中的某个节点
   *
   * @param value 要删除的节点的值
   */
  delete(value: T) {
    // 先找到要删除的节点
    const nodeToDelete = this.search(value);
    // 如果不存在，就直接退出
    if (!nodeToDelete) {
      return;
    }

    // 否则删除节点
    this._delete(nodeToDelete);
  }

  /**
   * 删除红黑树中的节点
   * @param node 要删除的节点
   */
  private _delete(node: RedBlackNode<T>) {
    // 如果该节点同时存在左右节点，则找到右子树的最小节点作为该节点的后继
    if (node.left && node.right) {
      const successor = this.minimum(node.right);
      node.value = successor!.value;
      node = successor!;
    }

    let child: RedBlackNode<T> | null;
    // 如果该节点存在左节点，则将该左节点作为它的唯一子节点
    if (node.left) {
      child = node.left;
    } else if (node.right) {
      // 如果该节点存在右节点，则将该右节点作为它的唯一子节点
      child = node.right;
    } else {
      child = null;
    }

    // 如果该节点没有子节点，直接删除
    if (!child) {
      // 如果该节点是黑色，则需要特殊处理
      if (node.color === Color.BLACK) {
        this._deleteCase1(node);
      }
      this._removeNode(node);
    } else {
      // 如果该节点是黑色，则需要特殊处理
      if (node.color === Color.BLACK) {
        // 如果该节点的唯一子节点是红色，则将该唯一子节点设置为黑色
        if (child.color === Color.RED) {
          child.color = Color.BLACK;
        } else {
          this._deleteCase1(node);
        }
      }
      // 用该节点的唯一子节点替换该节点
      this._replaceNode(node, child);
    }
  }

  private _deleteCase1(node: RedBlackNode<T>) {
    // 如果有父节点，就进入 Case 2
    if (node.parent) {
      this._deleteCase2(node);
    }
  }

  private _deleteCase2(node: RedBlackNode<T>) {
    // 找到兄弟节点
    const sibling = this._sibling(node);
    // 如果兄弟节点存在且颜色为红色
    if (sibling && sibling.color === Color.RED) {
      // 父节点颜色变为红色
      node.parent!.color = Color.RED;
      // 兄弟节点颜色变为黑色
      sibling.color = Color.BLACK;
      // 如果删除的节点是左子节点
      if (node === node.parent!.left) {
        // 则向左旋转
        this.leftRotate(node.parent!);
      } else {
        // 否则向右旋转
        this.rightRotate(node.parent!);
      }
    }
    this._deleteCase3(node);
  }


  private _deleteCase3(node: RedBlackNode<T>) {
    const sibling = this._sibling(node);
    // 当父节点颜色是黑色，兄弟节点颜色是黑色，兄弟节点的左右子节点都是黑色
    if (
      node.parent!.color === Color.BLACK &&
      sibling &&
      sibling.color === Color.BLACK &&
      (!sibling.left || sibling.left.color === Color.BLACK) &&
      (!sibling.right || sibling.right.color === Color.BLACK)
    ) {
      // 将兄弟节点颜色设置为红色
      sibling.color = Color.RED;
      // 递归处理父节点
      this._deleteCase1(node.parent!);
    } else {
      // 进入下一个情况
      this._deleteCase4(node);
    }
  }

  private _deleteCase4(node: RedBlackNode<T>) {
    const sibling = this._sibling(node);
    // 当父节点为红色，兄弟节点为黑色，且兄弟节点的左右子树为黑色时
    if (
      node.parent!.color === Color.RED &&
      sibling &&
      sibling.color === Color.BLACK &&
      (!sibling.left || sibling.left.color === Color.BLACK) &&
      (!sibling.right || sibling.right.color === Color.BLACK)
    ) {
      // 将兄弟节点涂红色
      sibling.color = Color.RED;
      // 父节点涂黑色
      node.parent!.color = Color.BLACK;
    } else {
      // 否则进入下一个删除 case
      this._deleteCase5(node);
    }
  }  

  private _deleteCase5(node: RedBlackNode<T>) {
    const sibling = this._sibling(node);
    if (sibling && sibling.color === Color.BLACK) {
      // 如果当前节点是它父的左节点，并且兄弟节点的右节点存在且为红色
      if (
        node === node.parent!.left &&
        sibling.right &&
        sibling.right.color === Color.RED
      ) {
        // 将兄弟节点的颜色设置为红色
        sibling.color = Color.RED;
        // 兄弟节点的右节点设置为黑色
        sibling.right!.color = Color.BLACK;
        // 对兄弟节点进行左旋
        this.leftRotate(sibling);
      } else if (
        node === node.parent!.right &&
        sibling.left &&
        sibling.left.color === Color.RED
      ) {
        // 同上
        sibling.color = Color.RED;
        sibling.left!.color = Color.BLACK;
        this.rightRotate(sibling);
      }
    }
    this._deleteCase6(node);
  }


  private _deleteCase6(node: RedBlackNode<T>) {
    const sibling = this._sibling(node);
    // 将兄弟节点颜色设置成父节点颜色
    sibling!.color = node.parent!.color;
    // 将父节点颜色设置成黑色
    node.parent!.color = Color.BLACK;
    if (node === node.parent!.left) {
      // 将兄弟节点的右子节点颜色设置成黑色
      sibling!.right!.color = Color.BLACK;
      // 对父节点左旋
      this.leftRotate(node.parent!);
    } else {
      // 将兄弟节点的左子节点颜色设置成黑色
      sibling!.left!.color = Color.BLACK;
      // 对父节点右旋
      this.rightRotate(node.parent!);
    }
  }

  private _removeNode(node: RedBlackNode<T>) {
    if (!node.parent) {
      this.root = null;
    } else if (node === node.parent.left) {
      node.parent.left = null;
    } else {
      node.parent.right = null;
    }
  }

  private _replaceNode(oldNode: RedBlackNode<T>, newNode: RedBlackNode<T>) {
    if (!oldNode.parent) {
      this.root = newNode;
    } else if (oldNode === oldNode.parent.left) {
      oldNode.parent.left = newNode;
    } else {
      oldNode.parent.right = newNode;
    }
    newNode.parent = oldNode.parent;
  }

  private _sibling(node: RedBlackNode<T>) {
    if (!node.parent) {
      return null;
    }
    return node === node.parent.left ? node.parent.right : node.parent.left;
  }
}

const rbtree = new RedBlackTree<number>();
rbtree.insert(10);
rbtree.insert(20);
rbtree.insert(30);
rbtree.insert(40);
rbtree.insert(50);
rbtree.insert(60);
rbtree.insert(70);
rbtree.insert(80);

rbtree.delete(10);
rbtree.delete(30);
rbtree.delete(20);
btPrint(rbtree.root);

export {};
